# 自定义数据集类
from torch.utils.data import DataLoader,Dataset
import torch
# class MultimodalDataset(Dataset):
#     def __init__(self):
#         self.images = [torch.randn(3, 224, 224) for _ in range(5)]
#         self.texts = [torch.randint(0, 10, (10,)) for _ in range(5)]

#     def __len__(self):
#         return len(self.images)

#     def __getitem__(self, idx):
#         return self.images[idx], self.texts[idx]
    
# dataset = MultimodalDataset()

# # 自定义 collate_fn 函数
# def multimodal_collate_fn(batch):
#     '''
    
#     batch 就是 batch_size=3 之后形成的[dataset,dataset,dataset]  
#     '''
#     print(batch)
    
#     images = [item[0] for item in batch]
#     texts = [item[1] for item in batch]
#     # 将图像组合成一个批量张量
#     image_batch = torch.stack(images)
#     # 将文本组合成一个批量张量
#     text_batch = torch.stack(texts)
#     return image_batch, text_batch

# # 创建数据集对象
# dataset = MultimodalDataset()

# # 创建 DataLoader 对象，指定 collate_fn 函数
# dataloader = DataLoader(dataset, batch_size=3, shuffle=False, collate_fn=multimodal_collate_fn)

# # 遍历 DataLoader
# for image_batch, text_batch in dataloader:
#     print("Image batch shape:", image_batch.shape)
#     print("Text batch shape:", text_batch.shape)





dataset = [(torch.randn(size=(5,)),3) for _ in range(10)]
# dataset = [(np.random.randn(12,),3) for _ in range(10)]

dl = DataLoader(dataset,batch_size=3,shuffle=False)
for _ in dl:
    print(_)
    # break

'''
[tensor([[-0.8831,  1.0164, -1.3281, -0.7386, -0.3061],
        [-0.6205, -0.0295,  0.1605,  0.7687,  0.6787],
        [-0.1642, -1.1912, -0.0672,  0.9157, -0.5048]]), tensor([3, 3, 3])]
[tensor([[ 1.3091, -0.8015,  0.1910, -0.9297, -0.4911],
        [ 2.0757, -0.2965,  2.4178,  1.6196, -1.2657],
        [ 0.5237, -0.6205, -1.3344,  1.4578, -0.4696]]), tensor([3, 3, 3])]
[tensor([[-0.2797,  2.0807, -1.3000, -2.3417,  0.0430],
        [-1.6221,  1.0319,  2.0326,  0.1292,  0.9023],
        [ 0.5434,  0.8136,  0.6429,  0.2411,  1.5817]]), tensor([3, 3, 3])]
[tensor([[ 1.0925, -0.4716,  0.3787, -1.4387, -1.8981]]), tensor([3])]


'''



def change(batchs):
    # batch 就是上面输出的那个
    '''
    [(tensor([ 0.9818,  0.1940,  0.0436, -0.8141,  1.4636]), 3), (tensor([-0.0840, -0.0065, -2.6674, -0.3946,  2.8424]), 3), (tensor([ 2.2204, -1.0966, -0.2777,  0.3998,  1.1074]), 3)]
    '''
    print(batchs)
    data=[]
    label = []
    for batch in batchs:
        data.extend(batch[0])
        label.append(batch[1])
    print('data',data)
    return  torch.tensor(data),torch.tensor(label)

dl = DataLoader(dataset,batch_size=3,shuffle=False,collate_fn=change)
for _ in dl:
    print(_)
    break
'''
(tensor([-0.8561,  0.3045, -0.6356,  0.6673,  0.6331, -0.0856,  0.8673, -0.7812,
        -1.7579,  1.4737, -0.1817,  0.4981,  1.4502,  1.3990,  0.4049]), tensor([3, 3, 3]))

'''
    
